//
// Created by hamlet on 23-3-11.
//

#include "Model.h"
#include "preprocess.h"

void Model::clear() {
    tmp_label.clear();
    tmp_word_index.clear();
    tmp_word_times.clear();
    tmp_file_word_times.clear();
    tmp_filtered_word_index.clear();
}

FansionML::IndexT Model::newLabel(const char *name) {
    auto i = label_name.size();
    label_name.emplace_back(name);
    return i;
}

FansionML::IndexT Model::newFileData(const char *name, FansionML::IndexT label_index) {
    auto i = file_name.size();
    file_name.emplace_back(name);
    tmp_label.emplace_back(label_index);
    tmp_file_word_times.emplace_back();
    return i;
}

void Model::addWord(const std::string &word, unsigned long file_index) {
    auto &tfwti = tmp_file_word_times[file_index];
    auto it1 = tmp_word_index.find(word);
    if (it1 == tmp_word_index.end()) {
        auto i = word_name.size();
        word_name.emplace_back(word);
        tmp_word_index.emplace(word, i);
        tmp_word_times.emplace_back(1);
        tmp_file_word_times[file_index].emplace(i, 1);
    } else {
        auto i = it1->second;
        ++tmp_word_times[i];
        auto it2 = tfwti.find(i);
        if (it2 == tfwti.end()) {
            tfwti.emplace(i, 1);
        } else {
            ++it2->second;
        }
    }
}

void Model::filterWords(int times) {
    FansionML::IndexT bias = 0;
    for (FansionML::IndexT i = 0; i < word_name.size(); ++i) {
        if (tmp_word_times[i] < times) {
            ++bias;
        } else {
            word_name[i - bias] = word_name[i];
            tmp_filtered_word_index.emplace(i, i - bias);
        }
    }
    word_name.resize(tmp_filtered_word_index.size());
}

void Model::transformToEigen() {
    std::vector<Eigen::Triplet<double>> tripletList;
    tripletList.reserve(file_name.size() * 12);
    for (FansionML::IndexT file_index = 0; file_index < file_name.size(); ++file_index) {
        for (auto &p: tmp_file_word_times[file_index]) {
            auto it = tmp_filtered_word_index.find(p.first);
            if (it != tmp_filtered_word_index.end()) {
                tripletList.emplace_back(file_index, it->second, p.second);
            }
        }
    }
    data = FansionML::DataT{static_cast<Eigen::Index>(file_name.size()),
                            static_cast<Eigen::Index>(tmp_filtered_word_index.size())};
    data.setFromTriplets(tripletList.begin(), tripletList.end());
    label = FansionML::LabelT{tmp_label.size()};
    for (FansionML::IndexT i = 0; i < tmp_label.size(); ++i) {
        label(static_cast<Eigen::Index>(i)) = static_cast<int>(tmp_label[i]);
    }
}

void Model::preprocess(const char *path) {
    ::preprocess(this, path);
}

void Model::calcTfPlusIdf() {
    for (int k = 0; k < data.cols(); ++k) {
        auto col = data.col(k);
        double idf = log(static_cast<double>(data.cols()) / col.sum());  // idf = log(N/n_k)
        col *= idf;  // a_ik = f_ik * idf
    }
}
